import os
import jieba
import datetime

def get_files_list(file_dir, postfix):
    File_Name=[]
    for files in os.listdir(file_dir):
        if os.path.splitext(files)[1] == postfix:
            File_Name.append(file_dir + '/' + files)
    return File_Name

# 字词分割，对整个文件内容进行字词分割
def segment_lines(file_list,segment_out_dir,stopwords=[]):
    for i,file in enumerate(file_list):
        segment_out_name=os.path.join(segment_out_dir,'segment_{}.txt'.format(i))
        with open(file, 'rb') as f:
            document = f.read()
            document_cut = jieba.cut(document)
            sentence_segment=[]
            for word in document_cut:
                if word not in stopwords:
                    sentence_segment.append(word)
            result = ' '.join(sentence_segment)
            result = result.encode('utf-8')
            with open(segment_out_name, 'wb') as f2:
                f2.write(result)

def segment():
    # 对source中的txt文件进行分词，输出到segment目录中
    file_list=get_files_list('./source', postfix='.txt')
    segment_lines(file_list, './segment')

def train():
    # 将Word转换成Vec，然后计算相似度
    from gensim.models import word2vec
    import multiprocessing
    # 如果目录中有多个文件，可以使用PathLineSentences
    sentences = word2vec.PathLineSentences('./segment')
    # 设置模型参数，进行训练
    model = word2vec.Word2Vec(sentences, size=100, window=3, min_count=1)
    model.save('./models/word2Vec_s100_w3_mc1.model')

    # 设置模型参数，进行训练
    model2 = word2vec.Word2Vec(sentences, size=128, window=5, min_count=5, workers=multiprocessing.cpu_count())
    # 保存模型
    model2.save('./models/word2Vec_s128_w5_mc5.model')

def test():
    from gensim.models import Word2Vec
    model1 = Word2Vec.load('./models/word2Vec_s100_w3_mc1.model')
    print("model1:")
    print(model1.wv.most_similar('曹操'))
    print(model1.wv.most_similar(positive=['曹操', '刘备'], negative=['张飞']))

    print("model2:")
    model2 = Word2Vec.load('./models/word2Vec_s128_w5_mc5.model')
    print(model2.wv.most_similar('曹操'))
    print(model2.wv.most_similar(positive=['曹操', '刘备'], negative=['张飞']))
def main():
    #mode = "train"
    mode = "test"
    if mode == "train":
        train()
    else:
        test()
if __name__ == "__main__":
    start_time = datetime.datetime.now()
    main()
    end_time = datetime.datetime.now()
    print(start_time)
    print(end_time)

